import subprocess as sp
import socket
import time
import csv

from util.cpu import CPU
from util.gpu import GPU
from util.fps import SurfaceFlingerFPS
from util.comm import ClientSocket
from util.argparser import DeviceParser

if __name__=="__main__":
    #################
    # Get arguments #
    #################
    args = DeviceParser()
    app = args.app
    exp_time = args.exp_time
    pixel_ip = args.pixel_ip
    pixel_port = args.pixel_port
    server_ip = args.server_ip
    server_port = args.server_port

    #########################
    # Initialize the device #
    #########################
    ''' Connect wirelessly to the device '''
    if (pixel_port != None):
        sp.check_output(['adb', 'kill-server'])
        output = sp.check_output(['adb', 'connect', pixel_ip+':'+pixel_port])
        print(output.decode('utf-8'))
    
    ''' Set all governors to userspace '''
    CPU.set_scaling_governor(pixel_ip, 'userspace')
    # GPU.set_scaling_governor(pixel_ip, 'userspace')

    ''' Print all current frequencies '''
    # CPU.print_all_freq(pixel_ip)

    ''' Create CPU and GPU instances '''
    # CPU 0(little), 4(middle), 7(big)
    cpus = []
    for i in [0, 4, 7]:
        cpu = CPU(idx = i, ip = pixel_ip)
        cpus.append(cpu)
    gpu = GPU(ip = pixel_ip)
    
    ''' Set CPU and GPU frequencies to maximum '''
    # for cpu in cpus:
        # cpu.set_freq(0) # (cpu.max_freq_idx)
    # gpu.set_freq(gpu.max_freq_idx)

    ''' Start FPS recorder '''
    view = "\"SurfaceView - com.garena.game.codm/com.tencent.tmgp.cod.CODMainActivity#0\""
    fps_driver = SurfaceFlingerFPS(view, ip=pixel_ip)

    ''' Connect to the agent server '''
    socket = ClientSocket(server_ip, server_port)
    
    #########################
    # Start collecting data #
    #########################
    try:
        for i in range(exp_time):
            time.sleep(2)
            # Record processor frequencies
            for cpu in cpus:
                cpu.log_freq()
            gpu.log_freq()
            # Record processor temperature
            # CPU.get_overall_temp(pixel_ip)
            # GPU.get_overall_temp(pixel_ip)
            # Record FPS
            fps = float(fps_driver.getFPS())
            print(f'[FPS] {fps}')
            # Send the state to the agent
            cpu_little_freq = cpus[0].get_freq_idx()
            cpu_big_freq = cpus[1].get_freq_idx()
            temp = CPU.get_overall_temp(pixel_ip)
            state_msg = f'{cpu_little_freq},{cpu_big_freq},{temp},{fps}'
            # print(f'state_msg: {state_msg}')
            socket.send(state_msg)

            # Wait until get actions from the agent
            action = socket.receive()
            # TODO: set GPU freq
            for i in [0]:
                cpus[i].set_freq(action[0], action_max=3)
            for i in [1, 2]:
                cpus[i].set_freq(action[1], action_max=3)

        # Close socket
        socket.send('1')
        
    #################
    # Stop learning #
    #################
    finally:
        ''' Output data to a log file '''
        csv_file = f'./output/231106-device_{exp_time}.csv'
        with open(csv_file, mode = 'w', newline = '') as file:
            writer = csv.writer(file)
            # Record processor frequencies and temperature
            for cpu in cpus:
                writer.writerow(cpu.freq_data)
            writer.writerow(gpu.freq_data)
            writer.writerow(CPU.temp_data)
            # writer.writerow(GPU.temp_data)
            # TODO: Record power consumption
            # Record FPS
            writer.writerow(fps_driver.fps_data)

        ''' Restore the frequency governing method '''
        CPU.set_scaling_governor(pixel_ip, 'schedutil')
        GPU.set_scaling_governor(pixel_ip)
